#该脚本的目的是为了适配torchvision.datasets.ImageFolder类的文件夹组织格式
import argparse
from pathlib import Path
import os
import random
import shutil

from tqdm import tqdm

def Parse_Arguments():
    parser = argparse.ArgumentParser(description="")
    parser.add_argument('--origin-img-path', default=r"/mnt/data/ubuntu/dataset/clip", type=str)  #类别文件夹所处的目录
    parser.add_argument('--target-img-path', default=r"/mnt/data/ubuntu/dataset/clip_train", type=str)  #移动图片所在的目标目录
    return parser.parse_args()

def main():
    args = Parse_Arguments()
    origin_img_path = Path(args.origin_img_path)
    origin_img_dir = [Path(_.name) for _ in origin_img_path.iterdir() if _.is_dir()]
    target_img_path = Path(args.target_img_path)
    test_frac = 0.2
    for j in origin_img_dir:
        target_class_img_train_dir = target_img_path / Path("train") / j
        target_class_img_test_dir = target_img_path / Path("test") / j
        os.makedirs(str(target_class_img_train_dir), exist_ok=True)
        os.makedirs(str(target_class_img_test_dir), exist_ok=True)
        origin_class_img_dir = origin_img_path / j
        img_paths = os.listdir(str(origin_class_img_dir))
        random.shuffle(img_paths) # 随机打乱
        val_number = int(len(img_paths) * test_frac) # 测试集文件个数
        train_files = img_paths[val_number:]         # 训练集文件名列表
        val_files = img_paths[:val_number]           # 测试集文件名列表
        # #训练集
        for each in tqdm(train_files):
            src_path = os.path.join(str(origin_class_img_dir), each)
            dst_path = os.path.join(str(target_class_img_train_dir),each)
            shutil.copy(src_path, dst_path)
        # #测试集
        for each in tqdm(val_files):
            src_path = os.path.join(str(origin_class_img_dir), each)
            dst_path = os.path.join(str(target_class_img_test_dir),each)
            shutil.copy(src_path, dst_path)


if __name__ == "__main__":
    main()
